package com.zhangfb95.crosschain.server.service.impl.handler;

import com.zhangfb95.crosschain.server.infra.model.mapper.entity.ChainAccount;
import com.zhangfb95.crosschain.server.infra.util.EccCertUtil;
import com.zhangfb95.crosschain.server.service.SignHandler;
import jakarta.annotation.PostConstruct;
import org.bouncycastle.asn1.ASN1Encodable;
import org.bouncycastle.asn1.ASN1InputStream;
import org.bouncycastle.asn1.ASN1Integer;
import org.bouncycastle.asn1.ASN1Primitive;
import org.bouncycastle.asn1.ASN1Sequence;
import org.bouncycastle.asn1.DERSequenceGenerator;
import org.bouncycastle.asn1.x9.ECNamedCurveTable;
import org.bouncycastle.asn1.x9.X9ECParameters;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.springframework.stereotype.Component;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.math.BigInteger;
import java.security.PrivateKey;
import java.security.Security;
import java.security.Signature;
import java.security.interfaces.ECPrivateKey;

import static java.lang.String.format;

/**
 * @author zhangfubing
 * @since 2023/12/1
 */
@Component
public class EcdsaSignHandler implements SignHandler {

    @PostConstruct
    public void init() {
        Security.addProvider(new BouncyCastleProvider());
    }

    @Override
    public String supportedAlgType() {
        return "ecdsa";
    }

    @Override
    public byte[] sign(ChainAccount chainAccount, byte[] bytes) {
        PrivateKey privateKey = EccCertUtil.getPrivateKey(chainAccount.getPrivateKey());
        return ecdsaSignToBytes((ECPrivateKey) privateKey, bytes);
    }

    private byte[] ecdsaSignToBytes(ECPrivateKey privateKey, byte[] data) {
        try {
            X9ECParameters params = ECNamedCurveTable.getByName("secp256r1");
            BigInteger curveN = params.getN();

            Signature sig = Signature.getInstance("SHA256withECDSA", BouncyCastleProvider.PROVIDER_NAME);
            sig.initSign(privateKey);
            sig.update(data);
            byte[] signature = sig.sign();

            BigInteger[] sigs = preventMalleability(decodeECDSASignature(signature), curveN);

            try (ByteArrayOutputStream s = new ByteArrayOutputStream()) {
                DERSequenceGenerator seq = new DERSequenceGenerator(s);
                seq.addObject(new ASN1Integer(sigs[0]));
                seq.addObject(new ASN1Integer(sigs[1]));
                seq.close();
                return s.toByteArray();
            }
        } catch (Exception e) {
            throw new RuntimeException("Could not sign the message using private key", e);
        }
    }

    private static BigInteger[] decodeECDSASignature(byte[] signature) throws Exception {
        try (ByteArrayInputStream inStream = new ByteArrayInputStream(signature)) {
            ASN1InputStream asnInputStream = new ASN1InputStream(inStream);
            ASN1Primitive asn1 = asnInputStream.readObject();

            BigInteger[] sigs = new BigInteger[2];
            int count = 0;
            if (asn1 instanceof ASN1Sequence) {
                ASN1Sequence asn1Sequence = (ASN1Sequence) asn1;
                ASN1Encodable[] asn1Encodables = asn1Sequence.toArray();
                for (ASN1Encodable asn1Encodable : asn1Encodables) {
                    ASN1Primitive asn1Primitive = asn1Encodable.toASN1Primitive();
                    if (asn1Primitive instanceof ASN1Integer) {
                        ASN1Integer asn1Integer = (ASN1Integer) asn1Primitive;
                        BigInteger integer = asn1Integer.getValue();
                        if (count < 2) {
                            sigs[count] = integer;
                        }
                        count++;
                    }
                }
            }
            if (count != 2) {
                throw new RuntimeException(format("Invalid ECDSA signature. Expected count of 2 but got: %d.", count));
            }
            return sigs;
        }

    }

    private BigInteger[] preventMalleability(BigInteger[] sigs, BigInteger curveN) {
        BigInteger cmpVal = curveN.divide(BigInteger.valueOf(2L));

        BigInteger sval = sigs[1];

        if (sval.compareTo(cmpVal) > 0) {

            sigs[1] = curveN.subtract(sval);
        }

        return sigs;
    }
}
